import argparse
from cmath import e
from concurrent import futures
from functools import cached_property
import logging
from pathlib import Path
import re
from Bio import Entrez  # pylint: disable=E0401
from pydantic import BaseModel
import polars as pl

logger = logging.getLogger("Prepare Gene EntrezID")
# 设置 NCBI API 密钥（可选）
# Entrez.email = "your_email@example.com"


class PrepareGeneID(BaseModel):
    ncbi_refseq_curated_file: Path
    gene_id_file: Path

    @cached_property
    def genepred_records(self) -> pl.DataFrame:
        logger.info("Reading %s", self.ncbi_refseq_curated_file)
        records = pl.read_csv(self.ncbi_refseq_curated_file, separator="\t", has_header=False, columns=[2, 12], new_columns=["chromosome", "gene"])
        records = records.filter(pl.col("chromosome").str.len_chars().le(5)).unique(["chromosome", "gene"])
        return records

    def fetch_entrez_id(self, chromosome: str, gene: str) -> str | None:
        # 构建查询字符串
        chrom = re.sub(r"[Cc][Hh][Rr]", "", chromosome)
        query = f"{gene}[Gene] AND {chrom}[Chromosome]"
        record = Entrez.read(Entrez.esearch(db="gene", term=query))
        # 获取 Entrez ID
        entrez_ids = record.get("IdList")
        entrez_id = str(entrez_ids[0]) if entrez_ids else None
        print(chromosome, gene, entrez_id)
        return {"chromosome": chromosome, "gene": gene, "gene_id": entrez_id}

    def run(self) -> pl.DataFrame:
        with futures.ProcessPoolExecutor(max_workers=32) as executor:
            tasks = []
            for row in self.genepred_records.iter_rows(named=True):
                tasks.append(executor.submit(self.fetch_entrez_id, **row))
            results = []
            for task in futures.as_completed(tasks):
                results.append(task.result())
            pl.DataFrame(results).write_csv(self.gene_id_file, separator="\t")


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description="Prepare Refseq")
    parser.add_argument("--ncbi_refseq_curated_file", "-i", type=Path, required=True, help="input, ncbi refseq curated_file")
    parser.add_argument("--gene_id_file", "-o", type=Path, required=True, help="outdir, outdir")
    parser.set_defaults(func=lambda args: PrepareGeneID.model_validate(vars(args)).run())
    args = parser.parse_args()
    args.func(args)
